#include <ntstatus.h>
#define WIN32_NO_STATUS

#include "minicon.h"
#include "miniopt.h"
#include <shellapi.h>
#include <winternl.h>
#include <stierr.h>

extern int err_ret(int errcode, const char *msg, const WCHAR *name);
extern bool real_filename(const WCHAR *nm, WCHAR *out, DWORD ms, bool norm, int must_exist);

#if defined(_WIN64) || defined(_DEBUG) || !defined(_M_IX86)
static void __forceinline llmuladd(__int64 &sum,unsigned __int32 mul,unsigned __int32 add)
{ sum = (__int64)((unsigned __int64)sum * mul + add); }
static DWORD __forceinline lldiv( __int64 &dividend, DWORD divisor )
{ DWORD reminder = (DWORD)(dividend % divisor); dividend /= divisor; return reminder; }
#else
void __fastcall llmuladd(__int64 &sum,unsigned __int32 mul,unsigned __int32 add);
DWORD __fastcall lldiv( __int64 &dividend, DWORD divisor );
#endif

//===========================================================================================================
//===========================================================================================================

#define FileStreamInformation   (_FILE_INFORMATION_CLASS)22

typedef struct _FILE_STREAM_INFORMATION
{
	ULONG NextEntryOffset;
	ULONG StreamNameLength;
	LARGE_INTEGER StreamSize;
	LARGE_INTEGER StreamAllocationSize;
	WCHAR StreamName[1];
} FILE_STREAM_INFORMATION, *PFILE_STREAM_INFORMATION;

typedef DWORD (__stdcall *TFNNtQueryInformationFile)(
   IN HANDLE FileHandle,
	OUT PIO_STATUS_BLOCK IoStatusBlock,
	OUT PVOID FileInformation,
	IN ULONG Length,
	IN FILE_INFORMATION_CLASS FileInformationClass
);

typedef __out HANDLE (WINAPI *TFNFindFirstStreamW)(
	__in       LPCWSTR lpFileName,
	__in       STREAM_INFO_LEVELS InfoLevel,
	__out      LPVOID lpFindStreamData,
	__reserved DWORD dwFlags
);
typedef BOOL (APIENTRY *TFNFindNextStreamW)(
	__in  HANDLE hFindStream,
	__out LPVOID lpFindStreamData
);

//===========================================================================================================
//===========================================================================================================

struct streams_params
{
   int bsz1; union { void *pv; BYTE *pb; WCHAR *pw; } u1;
   int bsz2; union { void *pv; BYTE *pb; WCHAR *pw; } u2;

   bool recurse; bool recurse_symlink;
   bool skip_zone; bool skip_favicon; bool skip_encryptable;

   const WCHAR *stream_name;
   char stream_scan_type;
   char op;

   TFNFindFirstStreamW       first_stream;
   TFNFindNextStreamW        next_stream;
   TFNNtQueryInformationFile ntqi;

   __int64 c_total, c_dirs, c_links, c_found;
};
static streams_params *s_p = NULL;

//---

struct ff_streams {
   DWORD    stream_id;
   DWORD        attrs;
   LARGE_INTEGER size;
   const WCHAR  *name;
   int      name_size;

   union { HANDLE fh; HANDLE hf; } uh;
   LPVOID context;
   union { void *pv; WIN32_FIND_STREAM_DATA *fs; FILE_STREAM_INFORMATION *fsi; WIN32_STREAM_ID *fsid; } uf;
   int cb;
};

//===========================================================================================================

static void find_close_stream( ff_streams *ffs )
{
   if ( ffs->uh.fh == INVALID_HANDLE_VALUE )
      return;
   if ( s_p->stream_scan_type == 'f' )
      ::FindClose( ffs->uh.hf );
   else {
      DWORD dwr;
      if ( s_p->stream_scan_type == 'b' && ffs->context )
      { ::BackupRead(ffs->uh.hf, NULL, 0, &dwr, TRUE, FALSE, &ffs->context ); }
      ffs->context = NULL;
      ::CloseHandle( ffs->uh.fh );
   }
   ffs->uh.fh = INVALID_HANDLE_VALUE;
}

//---

static bool find_next_stream( ff_streams *ffs )
{
   if ( ffs->uh.fh == INVALID_HANDLE_VALUE )
      return false;
   switch ( s_p->stream_scan_type ) {
   case 'f':
      if ( !(*s_p->next_stream)(ffs->uh.hf, ffs->uf.fs) ) return false;
      ffs->name_size = ::lstrlenW(ffs->name = ffs->uf.fs->cStreamName);
      ffs->size = ffs->uf.fs->StreamSize;
      return true;
   case 'n':
      if ( ffs->uf.fsi->StreamNameLength == 0) return false;
      ffs->name_size = (int)(ffs->uf.fsi->StreamNameLength / 2);
      ffs->name = ffs->uf.fsi->StreamName;
      ffs->size = ffs->uf.fsi->StreamSize;
      if ( ffs->uf.fsi->NextEntryOffset == 0 ) ffs->uf.fsi->StreamNameLength = 0;
      else ffs->uf.fsi = (PFILE_STREAM_INFORMATION)((LPBYTE)ffs->uf.fsi + ffs->uf.fsi->NextEntryOffset);
      return true;
   case 'b': {
      int ok = -1;
      HANDLE hF = ffs->uh.fh; LPBYTE bf = (LPBYTE)ffs->uf.fsid;
      DWORD dw = (DWORD)offset_of(WIN32_STREAM_ID, cStreamName), dwr = 0;
      if ( ::BackupRead(hF, bf, dw, &dwr, FALSE, FALSE, &ffs->context) && dwr == dw ) {
         if ( (dw = ffs->uf.fsid->dwStreamNameSize) == 0 ) ok = 0;
         else {
            if ( dw <= (DWORD)(ffs->cb/2 - sizeof(WIN32_STREAM_ID)/sizeof(WCHAR)))
               if ( ::BackupRead(hF, bf+dwr, dw, &dwr, FALSE, FALSE, &ffs->context) && dwr == dw )
                  ok = 0;
         }
      }
      if ( ok == 0 ) {
         ffs->name_size = (int)(ffs->uf.fsid->dwStreamNameSize / 2);
         ffs->uf.fsid->cStreamName[ffs->name_size] = L'\0';
         ffs->name = ffs->uf.fsid->cStreamName;
         ffs->size = ffs->uf.fsid->Size;
         ffs->attrs = ffs->uf.fsid->dwStreamAttributes;
         ffs->stream_id = ffs->uf.fsid->dwStreamId;
         if ( ffs->uf.fsid->Size.QuadPart == 0 ) ok = +1;
         else {
            if (::BackupSeek(hF, ffs->size.LowPart, (DWORD)ffs->size.HighPart, &dw, &dwr, &ffs->context))
               ok = +1;
         }
      }
      if ( ok <= 0 ) find_close_stream( ffs );
      return ok >= 0; }
   default:
      return false;
   }
}

//---

static bool find_first_stream( const WCHAR *fnm, ff_streams *ffs, DWORD dwAttr )
{
   if ( s_p->stream_scan_type == 'f' ) {
      ffs->uh.hf = (*s_p->first_stream)(fnm, FindStreamInfoStandard, ffs->uf.fs, 0);
      if ( ffs->uh.hf == INVALID_HANDLE_VALUE ) return false;
      ffs->name_size = ::lstrlenW(ffs->name = ffs->uf.fs->cStreamName);
      ffs->size = ffs->uf.fs->StreamSize;
      return true;
   }
   else {
      ffs->uh.hf = ::CreateFileW(
         fnm, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 
         FILE_FLAG_BACKUP_SEMANTICS | ((dwAttr & FILE_ATTRIBUTE_REPARSE_POINT) ? FILE_FLAG_OPEN_REPARSE_POINT : 0), 
         NULL
      );
   }
   if ( INVALID_HANDLE_VALUE == ffs->uh.hf )
      return false;

   if ( s_p->stream_scan_type == 'n' ) {
      ffs->uf.fsi->StreamNameLength = ffs->uf.fsi->NextEntryOffset = 0;
      IO_STATUS_BLOCK ios;
      NTSTATUS status = (*s_p->ntqi)(ffs->uh.fh, &ios, ffs->uf.fsi, (DWORD)ffs->cb, FileStreamInformation);
      if ( status != STATUS_SUCCESS ) {
         ffs->uf.fsi->StreamNameLength = ffs->uf.fsi->NextEntryOffset = 0;
         return false;
      }
   }

   return find_next_stream( ffs );
}

//---

static const WCHAR data[] = L"::$DATA";
static const WCHAR zone[] = L":Zone.Identifier:$DATA";
static const WCHAR favi[] = L":favicon:$DATA";
static const WCHAR encr[] = L":encryptable:$DATA";

#define STREAM_IS(s) (ffs.name_size == idim(s)-1 && 0 == Win::bcmp(ffs.name,s,sizeof(s)-2))

static const char *stream_data_ids[] =
{ "INVALID"   // 0x00000000
, "DATA"      // 0x00000001
, "EA_DATA"   // 0x00000002
, "SECURITY"  // 0x00000003
, ""          // 0x00000004 // ALTERNATE
, "LINK"      // 0x00000005
, "PROPERTY"  // 0x00000006
, "OBJECT_ID" // 0x00000007
, "REPARSE"   // 0x00000008
, "SPARSE"    // 0x00000009
, "TXFS"      // 0x0000000a
};

//---

static bool check_found( const WIN32_FIND_DATAW *ff, Win::SearchContext *ctx );
static bool already_seen( const WCHAR *dir );

//---

static int enumerate_streams( const WCHAR *nm, DWORD dwAttr, bool allow_recurse )
{
   int ns = 0, nfree = s_p->bsz2 / (int)sizeof(WCHAR);
   WCHAR *pwo = s_p->u2.pw, fnm[MAX_PATH*2], *fnm_e;

   if ( INVALID_FILE_ATTRIBUTES == dwAttr )
      return err_ret(-41, "can't open:", nm);
   if (!::GetFullPathNameW(nm, (DWORD)dim(fnm), fnm, NULL))
      return err_ret(-42, "bad filename:", nm);

   *pwo = L'\0'; fnm_e = fnm + ::lstrlenW(fnm);

   ++s_p->c_total;
   DWORD dir_link = (dwAttr & (FILE_ATTRIBUTE_DIRECTORY | FILE_ATTRIBUTE_REPARSE_POINT));
   if (dir_link == FILE_ATTRIBUTE_DIRECTORY) ++s_p->c_dirs;
   else if (dir_link) ++s_p->c_links;

   ff_streams ffs;
   ffs.stream_id = BACKUP_ALTERNATE_DATA;
   ffs.attrs = INVALID_FILE_ATTRIBUTES;
   ffs.size.QuadPart = -1;
   ffs.name = NULL;
   ffs.name_size = -1;
   ffs.uh.fh = INVALID_HANDLE_VALUE; ffs.context = NULL;
   ffs.uf.pv = s_p->u1.pv;
   ffs.cb = s_p->bsz1;

   bool ok = find_first_stream(fnm, &ffs, dwAttr );
   while (ok) {
      const WCHAR *sname = s_p->stream_name;
      int nw = sname ? ::lstrlenW(sname) : 0;
      bool alt = (ffs.stream_id == BACKUP_ALTERNATE_DATA);
      //  ':' -- 1, stream_name ':$DATA' -- 6
      bool skip = ((s_p->op == 'd' || nw) && (!alt || ffs.name_size <= 1+6)) 
         || (nw && (1+nw+6 != ffs.name_size || 0 != Win::bcmp(sname, ffs.name+1, nw*2)))
         || (ffs.stream_id == BACKUP_DATA)                    // data stream
         || (alt && STREAM_IS(data))                          // also data stream
         || (s_p->skip_zone && alt && STREAM_IS(zone))        // ignored Zone.Identifier
         || (s_p->skip_favicon && alt && STREAM_IS(favi))     // ignored favicon
         || (s_p->skip_encryptable && alt && STREAM_IS(encr)) // ignored encryptable
      ;
      if ( !skip ) {
         ++ns;
         if ( s_p->op == 'd' ) { // delete
            if ( ffs.name_size-6 < nfree + 2 ) {
               Win::bmov(pwo, ffs.name, (ffs.name_size-6)*2);
               nfree -= ffs.name_size-6 + 1;
               pwo += ffs.name_size-6;
               *pwo++ = L'\0'; *pwo = L'\0';
            }
         }
         else { // enumerate
            if ( 1 == ns ) { Con::o(nm); Con::n(); }
            Con::o("  \""); Con::o(ffs.name, ffs.name_size);
            Con::o("\" size="); Con::o(ffs.size.QuadPart);
            if (INVALID_FILE_ATTRIBUTES != ffs.attrs) { Con::o(" attr="); Con::o((unsigned)ffs.attrs); }
            if ( BACKUP_ALTERNATE_DATA != ffs.stream_id ) {
               Con::o(" id=");
               if ( ffs.stream_id >= (DWORD)dim(stream_data_ids)) Con::o((unsigned)ffs.stream_id);
               else Con::o(stream_data_ids[ffs.stream_id]);
            }
            Con::n();
         }
      }
      ok = find_next_stream(&ffs);
   }
   find_close_stream(&ffs);
   if (ns > 0)
      ++s_p->c_found;

   // delete found stream(s)
   //
   if ( ns > 0 && s_p->op == 'd' ) {
      pwo = s_p->u2.pw;
      while (*pwo) {
         ::lstrcpyW(fnm_e, pwo);
         BOOL d_ok = ::DeleteFileW(fnm);
         Con::f("del \""); Con::o(fnm); Con::f("\" -- %s\r\n", d_ok ? "ok" : "failed");
         pwo += ::lstrlenW(pwo) + 1;
      }
   }

   // recursive scan (if required)
   //
   *fnm_e = L'\0';
   if ( allow_recurse && s_p->recurse && 0 != (dwAttr & FILE_ATTRIBUTE_DIRECTORY) ) {
      bool skip = (!s_p->recurse_symlink && 0 != (dwAttr & FILE_ATTRIBUTE_REPARSE_POINT))
         || (s_p->recurse_symlink && already_seen(fnm))
      ;
      if ( !skip ) {
         WCHAR spath[MAX_PATH*2]; Win::SearchContext ctx;
         Win::init_search_context(&ctx, spath,MAX_PATH*2, check_found, s_p, NULL);
         Win::search_files(&ctx, fnm, L"*", true);
      }
   }

   return ns;
}

//-----------------------------------------------------------------------------------------------------------

static bool check_found( const WIN32_FIND_DATAW *ff, Win::SearchContext *ctx )
{
   if (ctx->cb_path <= ctx->nb_path + ::lstrlenW(ff->cFileName))
      return false; // too long name

   ::lstrcpyW(ctx->path + ctx->nb_path, ff->cFileName);

   bool dir = (0 != (ff->dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY));
   bool rpt = (0 != (ff->dwFileAttributes & FILE_ATTRIBUTE_REPARSE_POINT));
   bool sym = s_p->recurse_symlink;
   ctx->ignore = (!sym && rpt) || (sym && dir && already_seen(ctx->path));

   enumerate_streams(ctx->path, ff->dwFileAttributes, false);
   ++ctx->found;
   return true;
}

//-----------------------------------------------------------------------------------------------------------

static bool hash_table_initialized = false;
static const int hash_table_size = 102001;
struct HashTableEntry
{
   __int64                 hash64;
   HashTableEntry *next_collision;
   int                name_length;
   WCHAR                  name[1];
};
static HashTableEntry **hash_table = NULL;

//---

static bool already_seen( const WCHAR *dir )
{
   if ( !hash_table_initialized ) {
      Win::alloc((void* &)hash_table, sizeof(HashTableEntry *) * hash_table_size);
      Win::bzero(hash_table, sizeof(HashTableEntry *) * hash_table_size);
      hash_table_initialized = true;
   }

   WCHAR real_name[512]; const WCHAR *pw, *real = real_name;
   if ( !real_filename(dir, real_name, (DWORD)dim(real_name), false, true ) )
      if ( !::GetFullPathNameW( dir, (DWORD)dim(real_name), real_name, NULL))
         real = dir;

   __int64 h64 = 0, t; int len = 0;
   for ( pw = real; *pw; ++pw ) { ++len; llmuladd(h64, 17, *pw); }
   t = h64; if (t < 0) t = -t; DWORD rem = lldiv(t, (DWORD)hash_table_size);
   HashTableEntry *ph = hash_table[rem];
   while ( ph ) {
      if ( h64 == ph->hash64 && 0 == ::lstrcmpW(real, ph->name) )
         return true;
      ph = ph->next_collision;
   }

   ph = (HashTableEntry *)Win::alloc(offset_of(HashTableEntry,name) + len*sizeof(WCHAR) + sizeof(WCHAR));
   if (!ph) { ::SetLastError(ERROR_NOT_ENOUGH_MEMORY); Win::panic("not enough memory", 90); }
   ph->hash64 = h64;
   ph->next_collision = hash_table[rem]; hash_table[rem] = ph;
   ph->name_length = len;
   ::lstrcpyW(ph->name, real);

   return false;
}

//-----------------------------------------------------------------------------------------------------------

int 
streams(char op, WCHAR **args,int na, int recurse, int skip_and_stat, const WCHAR *sname, char method)
{
   streams_params p; Win::bzero(s_p = &p, sizeof(p));
   p.bsz1 = 64 * 1024; Win::alloc(p.u1.pv, p.bsz1);
   p.bsz2 = 64 * 1024; Win::alloc(p.u2.pv, p.bsz2);

   p.op = op;
   p.stream_name = sname;

   p.recurse = recurse != 0;
   p.recurse_symlink = recurse > 0;

   p.skip_zone = (skip_and_stat & 0x01) != 0;
   p.skip_favicon = (skip_and_stat & 0x02) != 0;
   p.skip_encryptable = (skip_and_stat & 0x04) != 0;
   bool stat_out = (skip_and_stat & 0x10000) != 0;

   if ( method != 'b' && method != 'n') {
      (FARPROC&)p.first_stream = ::GetProcAddress(Win::hkern, "FindFirstStreamW");
      (FARPROC&)p.next_stream = ::GetProcAddress(Win::hkern, "FindNextStreamW");
      method = (NULL != p.first_stream && NULL != p.next_stream) ? 'f' : 'n';
   }
   if ( method == 'n' ) {
      (FARPROC&)p.ntqi = ::GetProcAddress(Win::hnt, "NtQueryInformationFile");
      if (NULL == p.ntqi) method = 'b';
   }
   p.stream_scan_type = method;

   int i, rc = 0;
   for (i = 0; i < na; ++i) {
      if ( !Opt::have_wildcard(args[i]) ) {
         DWORD dwAttr = ::GetFileAttributesW(args[i]);
         rc = enumerate_streams(args[i], dwAttr, true);
         if (rc < 0)
            break;
      }
      else {
         WCHAR srcn[512], *pw;
         WIN32_FIND_DATAW ff;
         HANDLE hFind = ::FindFirstFileW(args[i], &ff);
         if ( INVALID_HANDLE_VALUE == hFind ) continue; //NOT error
         do {
            if (Opt::is_dot(ff.cFileName))
               continue;

            ::lstrcpyW(srcn, args[i]);
            pw = srcn + ::lstrlenW(srcn);
            while ( pw > srcn && pw[-1] != L'\\' && pw[-1] != L'/' && pw[-1] != L':') --pw;
            ::lstrcpyW(pw, ff.cFileName);

            rc = enumerate_streams(srcn, ff.dwFileAttributes, true);
            if (rc < 0)
               break;

         } while ( ::FindNextFileW(hFind, &ff) );
         ::FindClose(hFind);
         if (rc < 0)
            break;
      }
   }
   if ( stat_out ) {
      Con::n();
      Con::o(  "total: "); Con::o(p.c_total);
      Con::o("  dirs: " ); Con::o(p.c_dirs );
      Con::o("  links: "); Con::o(p.c_links);
      Con::o("  found: "); Con::o(p.c_found);
      Con::n();
   }
   return rc;
}

//===========================================================================================================
//===========================================================================================================
